#!/bin/bash
set -e
# curl -sSfL https://gitee.com/sunnybug/pubshell/raw/main/root_tool/install_llvm21.sh | bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 错误处理函数
handle_error() {
    log_error "脚本执行失败，错误代码: $1"
    log_error "请检查错误信息并重试"
    exit 1
}

# 设置错误处理
trap 'handle_error $?' ERR

# 检查是否为root用户
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本需要root权限运行"
        log_info "请使用: sudo $0"
        exit 1
    fi
}

# 检查操作系统类型和版本
detect_os() {
    log_info "检测操作系统类型和版本..."

    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS_ID="$ID"
        OS_CODENAME="$VERSION_CODENAME"
        OS_NAME="$NAME"
        OS_VERSION="$VERSION_ID"
    else
        log_error "无法检测操作系统类型，缺少 /etc/os-release"
        exit 1
    fi

    if [[ "$OS_ID" == "debian" && "$OS_CODENAME" == "bookworm" ]]; then
        log_success "检测到Debian 12 (bookworm)，完全支持"
        LLVM_MIRROR_CODENAME="bookworm"
    elif [[ "$OS_ID" == "ubuntu" && "$OS_CODENAME" == "noble" ]]; then
        log_success "检测到Ubuntu 24.04 (noble)，完全支持"
        LLVM_MIRROR_CODENAME="noble"
    else
        log_warning "检测到的系统: $OS_NAME $OS_VERSION ($OS_CODENAME)"
        log_error "此脚本目前仅支持Debian 12 (bookworm) 和 Ubuntu 24.04 (noble)"
        exit 1
    fi
}

# 修复仓库问题
fix_repositories() {
    log_info "检查并修复仓库问题..."

    # 备份原始sources.list
    if [[ ! -f /etc/apt/sources.list.backup ]]; then
        cp /etc/apt/sources.list /etc/apt/sources.list.backup
        log_info "已备份原始sources.list"
    fi

    # 替换过期的阿里云security镜像为清华大学镜像（仅对debian有效）
    if [[ "$OS_ID" == "debian" ]]; then
        if grep -q "mirrors.aliyun.com/debian-security" /etc/apt/sources.list; then
            log_info "替换过期的阿里云security镜像..."
            sed -i 's|mirrors.aliyun.com/debian-security|mirrors.tuna.tsinghua.edu.cn/debian-security|g' /etc/apt/sources.list
        fi

        # 替换sources.list.d中的过期镜像
        for file in /etc/apt/sources.list.d/*.list; do
            if [[ -f "$file" ]] && grep -q "mirrors.aliyun.com/debian-security" "$file"; then
                log_info "修复文件: $file"
                sed -i 's|mirrors.aliyun.com/debian-security|mirrors.tuna.tsinghua.edu.cn/debian-security|g' "$file"
            fi
        done
    fi
}

# 更新包列表
update_package_lists() {
    log_info "更新包列表..."

    # 尝试更新，如果失败则重试
    local retry_count=0
    local max_retries=3

    while [[ $retry_count -lt $max_retries ]]; do
        if apt update; then
            log_success "包列表更新成功"
            return 0
        else
            retry_count=$((retry_count + 1))
            log_warning "包列表更新失败，尝试重试 ($retry_count/$max_retries)..."
            sleep 2
        fi
    done

    log_error "包列表更新失败，请检查网络连接和仓库配置"
    return 1
}

# 卸载旧版本LLVM/Clang
uninstall_old_versions() {
    log_info "检查旧版本LLVM/Clang..."

    # 获取所有已安装的llvm和clang包
    local all_packages=$(dpkg -l | grep -E "^ii\s+(llvm|clang)" | awk '{print $2}')

    # 过滤出非21版本的包
    local old_packages=""
    while IFS= read -r package; do
        # 只卸载非21版本的包
        # 保留所有包含-21的包
        if [[ ! "$package" =~ -21 ]]; then
            old_packages="$old_packages $package"
        fi
    done <<< "$all_packages"

    # 去除首尾空格
    old_packages=$(echo "$old_packages" | xargs)

    if [[ -n "$old_packages" ]]; then
        log_warning "发现旧版本包: $old_packages"
        log_info "正在卸载旧版本..."

        # 卸载旧版本包
        if apt remove --purge -y $old_packages; then
            log_success "旧版本卸载成功"
        else
            log_warning "部分旧版本卸载失败，继续安装..."
        fi

        # 清理不需要的包
        apt autoremove -y
    else
        log_info "未发现需要卸载的旧版本LLVM/Clang包"
    fi
}

# 添加LLVM仓库
add_llvm_repository() {
    log_info "添加LLVM仓库..."

    # 检查是否已经添加了LLVM仓库
    if [[ -f /etc/apt/sources.list.d/llvm.list ]]; then
        log_info "LLVM仓库已存在，跳过添加"
        return 0
    fi

    # 下载并添加GPG密钥
    log_info "下载LLVM GPG密钥..."
    if ! wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor > /etc/apt/trusted.gpg.d/llvm.gpg; then
        log_error "下载LLVM GPG密钥失败"
        return 1
    fi

    # 添加LLVM仓库
    log_info "添加LLVM仓库源..."
    if [[ "$OS_ID" == "debian" && "$LLVM_MIRROR_CODENAME" == "bookworm" ]]; then
        echo "deb http://mirrors.tuna.tsinghua.edu.cn/llvm-apt/bookworm llvm-toolchain-bookworm-21 main" | tee /etc/apt/sources.list.d/llvm.list
    elif [[ "$OS_ID" == "ubuntu" && "$LLVM_MIRROR_CODENAME" == "noble" ]]; then
        echo "deb http://mirrors.tuna.tsinghua.edu.cn/llvm-apt/noble llvm-toolchain-noble-21 main" | tee /etc/apt/sources.list.d/llvm.list
    else
        log_error "不支持的系统或版本，无法添加LLVM仓库"
        return 1
    fi

    log_success "LLVM仓库添加成功"
}

# 安装LLVM包
install_llvm_packages() {
    log_info "安装LLVM 21包..."

    local packages=(
        "llvm-21"
        "clang-21"
        "clang-format-21"
        "clang-tidy-21"
        "clang-tools-21"
        "libc++-21-dev"
        "libc++abi-21-dev"
        "clangd-21"
    )

    # 检查哪些包需要安装
    local packages_to_install=()
    for package in "${packages[@]}"; do
        if ! dpkg -l | grep -q "^ii\s\+$package\s"; then
            packages_to_install+=("$package")
        else
            log_info "包 $package 已安装，跳过"
        fi
    done

    if [[ ${#packages_to_install[@]} -eq 0 ]]; then
        log_info "所有LLVM包都已安装"
        return 0
    fi

    log_info "安装包: ${packages_to_install[*]}"

    # 尝试安装包
    local retry_count=0
    local max_retries=3

    while [[ $retry_count -lt $max_retries ]]; do
        if apt install --no-install-recommends -y "${packages_to_install[@]}"; then
            log_success "LLVM包安装成功"
            return 0
        else
            retry_count=$((retry_count + 1))
            log_warning "包安装失败，尝试重试 ($retry_count/$max_retries)..."
            sleep 5
        fi
    done

    log_error "LLVM包安装失败"
    return 1
}

# 创建符号链接
create_symbolic_links() {
    log_info "创建符号链接..."

    # 定义需要创建的符号链接
    local links=(
        "clang-21:clang"
        "clang++-21:clang++"
        "clang-format-21:clang-format"
        "clang-tidy-21:clang-tidy"
        "clang-apply-replacements-21:clang-apply-replacements"
        "clangd-21:clangd"
    )

    # 先删除可能存在的旧链接
    log_info "清理旧符号链接..."
    for link in "${links[@]}"; do
        local target="${link#*:}"
        if [[ -L "/usr/bin/$target" ]]; then
            rm -f "/usr/bin/$target"
            log_info "删除旧链接: /usr/bin/$target"
        fi
    done

    # 创建新的符号链接
    log_info "创建新符号链接..."
    for link in "${links[@]}"; do
        local source="${link%:*}"
        local target="${link#*:}"

        # 检查源文件是否存在
        if [[ ! -f "/usr/bin/$source" ]] && [[ ! -L "/usr/bin/$source" ]]; then
            log_warning "源文件不存在: /usr/bin/$source，跳过创建链接"
            continue
        fi

        if ln -sf "/usr/bin/$source" "/usr/bin/$target"; then
            log_info "创建链接: $target -> $source"
        else
            log_warning "创建链接失败: $target -> $source"
        fi
    done

    log_success "符号链接创建完成"
}

# 验证安装
verify_installation() {
    log_info "验证LLVM 21安装..."

    local tools=("clang" "clang++" "clangd" "clang-format" "clang-tidy")
    local all_success=true

    for tool in "${tools[@]}"; do
        if command -v "$tool" >/dev/null 2>&1; then
            local version=$("$tool" --version 2>/dev/null | head -1)
            log_success "$tool 可用: $version"
        else
            log_error "$tool 不可用"
            all_success=false
        fi
    done

    if [[ "$all_success" == "true" ]]; then
        log_success "LLVM 21安装验证成功！"
        return 0
    else
        log_error "LLVM 21安装验证失败"
        return 1
    fi
}

# 清理函数
cleanup() {
    log_info "清理临时文件..."
    # 可以在这里添加清理逻辑
}

# 主函数
main() {
    log_info "开始安装LLVM 21..."

    # 检查root权限
    check_root

    # 检查操作系统
    detect_os

    # 修复仓库问题
    fix_repositories

    # 更新包列表
    update_package_lists

    # 卸载旧版本
    uninstall_old_versions

    # 添加LLVM仓库
    add_llvm_repository

    # 再次更新包列表（因为添加了新仓库）
    update_package_lists

    # 安装LLVM包
    install_llvm_packages

    # 创建符号链接
    create_symbolic_links

    # 验证安装
    verify_installation

    # 清理
    cleanup

    log_success "LLVM 21安装完成！"
    log_info "现在可以使用以下命令："
    log_info "  clang --version"
    log_info "  clang++ --version"
    log_info "  clangd --version"
    log_info "  clang-format --version"
    log_info "  clang-tidy --version"
}

# 运行主函数
main "$@"